import numpy as np
# import matplotlib.pyplot as plt
import os
import struct
import common.Codec as codec


def read_file(file):
    head_size = 10
    print('the current file is {0:s}'.format(file))
    # load data header
    all_content = np.fromfile(file, dtype='uint8', count=-1)
    if all_content is None:
        print("error head is null")
        return
    current = 0
    content_len = len(all_content)
    while current < content_len:
        # small
        sample_size = codec.hex2uint32(all_content[current:current + 4])
        # print('sample_size is {0:d}'.format(sample_size))
        tag_code = codec.hex2uint16(all_content[current + 4:current + 6])
        pic_width = codec.hex2uint16(all_content[current + 6:current + 8])
        pic_height = codec.hex2uint16(all_content[current + 8:current + 10])
        current = current + 10
        image_size = pic_height * pic_width
        # validate
        if sample_size != head_size + image_size:
            print("error. decode error")
            return
        image = np.array(all_content[current:current + image_size]) \
            .reshape(pic_height, pic_width)
        current = current + image_size
        # show image
        # plt.figure()
        # plt.imshow(image, cmap='gray')
        # plt.show()
        yield image, tag_code


# load all file from a dir
def get_all_images(directory, tail):
    # get all files and dirs
    file_name_list = os.listdir(directory)
    image_list = []
    image_label = []
    array_count = []
    for file_name in file_name_list:
        if file_name.endswith(tail):
            file = os.path.join(directory, file_name)
            len = 0
            for image, tag_code in read_file(file):
                image_list.append(image / 255.0)
                tag_code_unicode = struct.pack('<H', tag_code).decode('gb2312')
                image_label.append(tag_code_unicode)
                len = len + 1
        array_count.append(len)
    return image_list, image_label, array_count
